{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "005b25a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "54eee8e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone = torchvision.models.resnet18(pretrained=True).to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "db4eb7f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 5\n",
    "backbone.fc = torch.nn.Linear(backbone.fc.in_features, num_classes).to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "40a2b3b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load checkpoint\n",
    "backbone.load_state_dict(torch.load(\"CSOL-Utilities-resnet18-800x600.pth\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "1d1dbfc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "width = 800\n",
    "height = 600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "e3890b8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dir = \"C:/Users/Silver/BaiduSyncdisk/CSOL 集成工具数据集/\"\n",
    "class RemoveAlpha:\n",
    "    def __call__(self, img):\n",
    "        if img.mode == 'RGBA':\n",
    "            img = img.convert('RGB')\n",
    "        return img\n",
    "import random\n",
    "transform = torchvision.transforms.Compose([\n",
    "    RemoveAlpha(),\n",
    "    torchvision.transforms.Pad(padding=50, fill=0, padding_mode='constant'),\n",
    "    torchvision.transforms.Resize((height, width)),\n",
    "    torchvision.transforms.RandomHorizontalFlip(p=0.5),\n",
    "    torchvision.transforms.RandomVerticalFlip(p=0.5),\n",
    "    torchvision.transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    torchvision.transforms.RandomErasing(p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3), value=0),\n",
    "])\n",
    "dataset = torchvision.datasets.ImageFolder(dataset_dir, transform=transform)\n",
    "test_set_ratio = 0.1\n",
    "train_size = int(len(dataset) * (1 - test_set_ratio))\n",
    "test_size = len(dataset) - train_size\n",
    "train_set, test_set = torch.utils.data.random_split(dataset, [train_size, test_size])\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=16, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "ab3d7661",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(backbone.parameters(), lr=0.0001)\n",
    "num_epochs = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "967491c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/50], Loss: 0.0039\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [2/50], Loss: 0.0031\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [3/50], Loss: 0.0012\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [4/50], Loss: 0.0033\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [5/50], Loss: 0.0013\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [6/50], Loss: 0.0011\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [7/50], Loss: 0.0017\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [8/50], Loss: 0.0023\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [9/50], Loss: 0.0020\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [10/50], Loss: 0.0014\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [11/50], Loss: 0.0007\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [12/50], Loss: 0.0006\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [13/50], Loss: 0.0010\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [14/50], Loss: 0.0006\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [15/50], Loss: 0.0003\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [16/50], Loss: 0.0005\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [17/50], Loss: 0.0004\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [18/50], Loss: 0.0004\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [19/50], Loss: 0.0014\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [20/50], Loss: 0.0003\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [21/50], Loss: 0.0003\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [22/50], Loss: 0.0003\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [23/50], Loss: 0.0002\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [24/50], Loss: 0.0040\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [25/50], Loss: 0.0014\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [26/50], Loss: 0.0039\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [27/50], Loss: 0.0035\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [28/50], Loss: 0.0117\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [29/50], Loss: 0.0052\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [30/50], Loss: 0.0013\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [31/50], Loss: 0.0011\n",
      "Test Accuracy: 100.00%\n",
      "Epoch [32/50], Loss: 0.0012\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [33/50], Loss: 0.0004\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [34/50], Loss: 0.0019\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [35/50], Loss: 0.0077\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [36/50], Loss: 0.0017\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [37/50], Loss: 0.0011\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [38/50], Loss: 0.0007\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [39/50], Loss: 0.0006\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [40/50], Loss: 0.0008\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [41/50], Loss: 0.0007\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [42/50], Loss: 0.0017\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [43/50], Loss: 0.0017\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [44/50], Loss: 0.0006\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [45/50], Loss: 0.0004\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [46/50], Loss: 0.0002\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [47/50], Loss: 0.0002\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [48/50], Loss: 0.0003\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [49/50], Loss: 0.0002\n",
      "Test Accuracy: 98.31%\n",
      "Epoch [50/50], Loss: 0.0003\n",
      "Test Accuracy: 98.31%\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    backbone.train()\n",
    "    running_loss = 0.0\n",
    "    for images, labels in train_loader:\n",
    "        images = images.to('cuda')\n",
    "        labels = labels.to('cuda')\n",
    "        optimizer.zero_grad()\n",
    "        outputs = backbone(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "    print(f\"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}\")\n",
    "    backbone.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images = images.to('cuda')\n",
    "            labels = labels.to('cuda')\n",
    "            outputs = backbone(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    acc = 100 * correct / total\n",
    "    print(f\"Test Accuracy: {acc:.2f}%\")\n",
    "    if epoch > 0 and epoch % 50 == 0:\n",
    "        torch.save(backbone.state_dict(), f\"CSOL-Utilities-ResNet18-800x600-epoch{epoch}.pth\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "074ef804",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存 .pth\n",
    "torch.save(backbone.state_dict(), \"CSOL-Utilities-ResNet18-800x600.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "4fa0d154",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Silver\\AppData\\Local\\Temp\\ipykernel_36116\\2390526172.py:3: DeprecationWarning: You are using the legacy TorchScript-based ONNX export. Starting in PyTorch 2.9, the new torch.export-based ONNX exporter will be the default. To switch now, set dynamo=True in torch.onnx.export. This new exporter supports features like exporting LLMs with DynamicCache. We encourage you to try it and share feedback to help improve the experience. Learn more about the new export logic: https://pytorch.org/docs/stable/onnx_dynamo.html. For exporting control flow: https://pytorch.org/tutorials/beginner/onnx/export_control_flow_model_to_onnx_tutorial.html.\n",
      "  torch.onnx.export(\n"
     ]
    }
   ],
   "source": [
    "# 导出 .onnx\n",
    "dummy_input = (torch.randn(1, 3, height, width),)\n",
    "torch.onnx.export(\n",
    "    backbone.to('cpu'),\n",
    "    dummy_input,\n",
    "    \"CSOL-Utilities-ResNet18-800x600.onnx\",\n",
    "    export_params=True, # Store trained parameters\n",
    "    opset_version=11, # ONNX version\n",
    "    do_constant_folding=True, # Optimize constant folding\n",
    "    input_names=[\"screenshot\"],\n",
    "    output_names=[\"interface_type\"],\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
